#include "File/ChainLoader.h"
#include "Functions/Draw.h"
#include "Functions/Hist.h"
#include "Globals/Settings.h"
#include "TList.h"
#include "TTimeStamp.h"
#include <iomanip>
#include <iostream>
using namespace BOSS_Afterburner;
using namespace std;

/// Automatically book addresses for all branches of the underlying `TChain`. The branches are accessible by name (use method `Get_<type>`, though you have to the `type` to use this).
void ChainLoader::BookAddresses()
{
  if(!GetNbranches()) return;
  /// -# Loop over list of `TBranches` and determine `typename` of the branch. The data type of a branch can be determined from the last character of its title. See (`TTree`)["https://root.cern.ch/doc/master/classTTree.html#a8a2b55624f48451d7ab0fc3c70bfe8d7"] for the labels of each type.
  TIter next(GetListOfBranches());
  while(auto obj = next())
  {
    /// If branch title contains a `[`, this means it is an array and it will not be loaded.
    string type(obj->GetTitle());
    if(type.find('[') != string::npos) continue;
    switch(type.back())
    {
      case 'B': SetAddress(obj, fMap_B); break;
      case 'b': SetAddress(obj, fMap_b); break;
      case 'S': SetAddress(obj, fMap_S); break;
      case 's': SetAddress(obj, fMap_s); break;
      case 'I': SetAddress(obj, fMap_I); break;
      case 'i': SetAddress(obj, fMap_i); break;
      case 'F': SetAddress(obj, fMap_F); break;
      case 'D': SetAddress(obj, fMap_D); break;
      case 'L': SetAddress(obj, fMap_L); break;
      case 'l': SetAddress(obj, fMap_l); break;
      case 'O': SetAddress(obj, fMap_O); break;
      default: cout << "ERROR: Unable to book address for branch \"" << type << "\"" << endl;
    }
  }
}

void ChainLoader::Print(bool averages)
{
  /// -# Check if contains branches.
  if(!GetNbranches())
  {
    cout << "Chain \"" << GetName() << "\" has no branches." << endl;
    return;
  }

  /// -# Get list of `TBranches`.
  cout << "Tree \"" << GetName() << "\" (" << GetEntries() << " entries)" << endl;
  cout << GetTitle() << endl;
  cout << "   " << setw(18) << left << "BRANCH NAME";
  if(averages) cout << setw(12) << right << "MEAN" << endl;
  cout << endl;

  /// -# Loop over list of `TBranches` and print them.
  TIter next(GetListOfBranches());
  while(auto obj = next())
  {
    string title(obj->GetTitle());
    cout << "   " << setw(18) << left << title;
    if(averages) cout << setw(12) << right << ComputeMean(this, obj->GetName());
    if(title.find('[') != string::npos) cout << " (not loaded)" << endl;
  }
  cout << endl;
}

/// Draw a distribution of one of the branches in the file.
TH1F* ChainLoader::DrawHistogram(const char* branchX, const Int_t nBinx, const double x1,
                                 const double x2, Option_t* option, const bool save,
                                 const TString& logScale, const char* cut)
{
  // * Draw histogram * //
  const char* histName = Form("hist_%s", branchX);
  DrawHistogram(Form("%s>>%s(%i,%f,%f)", branchX, histName, nBinx, x1, x2), cut, option, save,
                logScale);
  // * Modify histogram * //
  TH1F* hist = (TH1F*)gDirectory->Get(histName);
  hist->SetTitle(Form("#it{%s} branch of the \"%s\" tree", branchX, GetName()));
  if(!logScale.Contains("y") && !logScale.Contains("Y"))
  { // NOTE: cannot set y log scale if min value is 0
    auto maxbin = hist->GetMaximumBin();
    hist->GetYaxis()->SetRangeUser(0.,
                                   hist->GetBinContent(maxbin) + 1.5 * hist->GetBinErrorUp(maxbin));
  }
  Hist::SetAxisTitles(hist, branchX, Form("count / %g", hist->GetBinWidth(1)));
  // * Save and return histogram * //
  if(save) Draw::SaveCanvas(Form("%s/%s", GetName(), branchX), gPad, logScale);
  return hist;
}

/// Draw a distribution of one of the branches in the file.
TH1* ChainLoader::DrawHistogram(const char* varexp, const char* cut, Option_t* option,
                                const bool save, const TString& logScale)
{
  /// <ol>
  /// <li> Check if `TChain` contains entries.
  if(!GetEntries()) return nullptr;
  /// <li> Remove possible output options from `varexp` parameter.
  TString varExp{varexp};
  TString outputName{varExp};
  TString histName{varExp};
  if(varExp.Contains(">>"))
  {
    auto firstPipe = outputName.First('>');
    outputName.Resize(firstPipe);
    histName.Remove(0, firstPipe + 2);
    histName.Strip(TString::kBoth, ' ');
  }
  else
  {
    TTimeStamp stamp;
    histName = Form("hist_%d", stamp.GetNanoSec());
    varExp += Form(">>%s", histName.Data());
  }
  /// <li> Draw histogram and save
  Draw(varExp.Data(), cut, option);
  auto hist = dynamic_cast<TH1*>(gDirectory->Get(histName.Data()));
  if(save) Draw::SaveCanvas(Form("%s/%s", GetName(), outputName.Data()), gPad, logScale);
  /// </ol>
  return hist;
}

/// Draw a distribution of one of the branches in the file.
TH2F* ChainLoader::DrawHistogram(const char* branchX, const char* branchY, const Int_t nBinx,
                                 const double x1, const double x2, const Int_t nBiny,
                                 const double y1, const double y2, Option_t* option,
                                 const bool save, const TString& logScale, const char* cut)
{
  // * Draw histogram * //
  const char* histName = Form("hist_%s_%s", branchX, branchY);
  DrawHistogram(
    Form("%s:%s>>%s(%i,%f,%f,%i,%f,%f)", branchX, branchY, histName, nBinx, x1, x2, nBiny, y1, y2),
    cut, option, save, logScale);
  // * Modify histogram * //
  TH2F* hist = (TH2F*)gDirectory->Get(histName);
  hist->SetTitle(Form("\"%s\" tree: #it{%s} vs #it{%s}", GetName(), branchX, branchY));
  Hist::SetAxisTitles(
    hist, branchX, branchY,
    Form("count / (%gx%g)", hist->GetXaxis()->GetBinWidth(1), hist->GetYaxis()->GetBinWidth(1)));
  // * Save and return histogram * //
  if(save) Draw::SaveCanvas(Form("%s/%s:%s", GetName(), branchY, branchX), gPad, logScale);
  return hist;
}

/// Draw a distribution of one of the branches in the file.
TH1* ChainLoader::DrawHistogram(const BranchPlotOptions& options)
{
  /// -# **Abort** if `TChain` contains entries.
  if(!GetEntries())
  {
    TerminalIO::PrintWarning(Form("Chain \"%s\" has no entries", GetName()));
    return nullptr;
  }
  /// -# Draw histogram and save
  Draw(options.VarExp(), options.CutSelection(), options.DrawOption());
  auto hist = dynamic_cast<TH1*>(gDirectory->Get(options.BuildHistName().c_str()));
  if(!hist)
  {
    TerminalIO::PrintWarning(Form("Cannot draw branch %s in tree %s", options.VarExp(), GetName()));
    return nullptr;
  }
  options.SetAxisLabels(hist);
  options.SetPlotLabel(hist);
  if(options.IsWrite())
    Draw::SaveCanvas(Form("%s/%s", GetName(), options.OutputFileName()), gPad, options.LogXYZ());
  return hist;
}

TH1* ChainLoader::DrawHistogram(const YAML::Node& input)
{
  if(!input) return nullptr;
  BranchPlotOptions options(input);
  options.SetTreeName(GetName());
  return DrawHistogram(options);
}

/// Draw the distributions of all branches of the underlying `TChain`.
void ChainLoader::DrawAndSaveAllBranches(Option_t* option, const TString& logScale)
{
  if(!GetNtrees()) return;
  TIter    next(GetListOfBranches());
  TObject* obj = nullptr;
  while((obj = next()))
    DrawHistogram(obj->GetName(), "", option, true, logScale);
}

/// Draw the distributions of all branches of the underlying `TChain` if its name starts with `"mult"`. This function additionally ensures that the bin width is set to 1 and that the histograms are drawn in `"E1"` mode (see [`THistPainter`](https://root.cern.ch/doc/master/classTHistPainter.html).
void ChainLoader::DrawAndSaveAllMultiplicityBranches(const TString& logScale, Option_t* option)
{
  if(!GetNtrees()) return;
  TString name(GetName());
  if(!name.BeginsWith("mult")) return;
  TIter    next(GetListOfBranches());
  TObject* obj = nullptr;
  while((obj = next()))
  {
    const char* name = obj->GetName();
    int         max  = GetMaximum(name);
    if(!max) continue;
    DrawHistogram(name, max, 0, max, option, true, logScale);
  }
}

/// Create a histogram object especially for invariant mass analysis.
TH1F* ChainLoader::GetInvariantMassHistogram(const char*                  varexp,
                                             const ReconstructedParticle& particle, const int nBins,
                                             Option_t* option, const TString& logScale)
{
  // * Check input arguments * //
  if(*varexp != 'm')
  {
    cout << "ERROR: varexp \"" << varexp << "\" does not start with 'm'" << endl;
    return nullptr;
  }
  // * Get histogram and modify
  auto hist = DrawHistogram(varexp, nBins, particle.PlotFrom(), particle.PlotUntil(), option, false,
                            logScale);
  hist->SetTitle(Form("Invariant mass for %s candidate", particle.NameLaTeX()));
  Hist::SetAxisTitles(hist, Form("#it{M}_{%s} (GeV/#it{c}^{2})", particle.DaughterLabel()),
                      Form("count / %g", hist->GetYaxis()->GetBinWidth(1)));
  return hist;
}

/// Create a histogram object based on a `BranchPlotOptions` and `ReconstructedParticle` object.
TH1F* ChainLoader::GetInvariantMassHistogram(const BranchPlotOptions&     branch,
                                             const ReconstructedParticle& particle)
{
  if(!branch.IsOK()) return nullptr;
  if(branch.GetNBranches() != 1)
  {
    TerminalIO::PrintWarning(Form("Cannot get an invariant mass plot for BranchPlotOptions "
                                  "object\n  \"%s\"\nbecause it has %u branches",
                                  branch.VarExp(), branch.GetNBranches()));
    return nullptr;
  }
  // * Check input arguments * //
  if(*branch.VarExp() != 'm')
  {
    cout << "ERROR: varexp \"" << branch.VarExp() << "\" does not start with 'm'" << endl;
    return nullptr;
  }
  // * Get histogram and modify
  auto hist = dynamic_cast<TH1F*>(DrawHistogram(branch));
  gPad->SetLogx(branch.LogX());
  gPad->SetLogy(branch.LogY());
  gPad->SetLogz(branch.LogZ());
  hist->SetTitle(Form("Invariant mass for %s candidate", particle.NameLaTeX()));
  Hist::SetAxisTitles(hist, Form("#it{M}_{%s} (GeV/#it{c}^{2})", particle.DaughterLabel()),
                      Form("count / %g", hist->GetYaxis()->GetBinWidth(1)));
  return hist;
}

/// Compute mean value of a branch in the `TChain`.
double ChainLoader::ComputeMean(TChain* chain, const char* branchName)
{
  if(!chain) return -1e6;
  chain->Draw(branchName);
  TH1F* htemp = (TH1F*)gPad->GetPrimitive("htemp");
  if(!htemp) return -1e6;
  double mean = htemp->GetMean();
  delete htemp;
  return mean;
}